Skip to content

[SPARK-58078][CORE] Add PipelinedShuffleDependency#57179

Open
jerrypeng wants to merge 1 commit into
apache:masterfrom
jerrypeng:stack/pipelined-shuffle-pr1-dependency
Open

[SPARK-58078][CORE] Add PipelinedShuffleDependency#57179
jerrypeng wants to merge 1 commit into
apache:masterfrom
jerrypeng:stack/pipelined-shuffle-pr1-dependency

Conversation

@jerrypeng

Copy link
Copy Markdown
Contributor

What changes were proposed in this pull request?

This PR adds PipelinedShuffleDependency, a new ShuffleDependency subtype that marks a
shuffle's output as incrementally readable: a consumer stage may begin reading the output
while the producer stage is still running, rather than waiting for the producer's full,
materialized output.

Class hierarchy:

Dependency[T]
ShuffleDependency[K, V, C]
PipelinedShuffleDependency[K, V, C] <- new

This PR adds only the type. On its own, PipelinedShuffleDependency behaves exactly like its
parent ShuffleDependency — construction goes through the normal path (shuffle id allocation,
ShuffleManager.registerShuffle, cleaner registration), and any code that matches on
ShuffleDependency continues to treat it as an ordinary, fully-materialized shuffle. It introduces
no behavior change. It is the marker that a later DAGScheduler change will use to (a)
co-schedule the producer and consumer stages connected by this edge (a "pipelined group") and (b)
select an incremental shuffle implementation for this edge. That concurrent-scheduling and
incremental-shuffle behavior is added in follow-up PRs.

The parent's checksumMismatchFullRetryEnabled / checksumMismatchQueryLevelRollbackEnabled
constructor parameters are intentionally not exposed by the subclass, so they stay at their false
defaults. Their checksum-mismatch model recomputes and re-runs succeeding stages after a mismatch;
in the pipelined-group failure model any failure aborts the whole group and the caller reruns from
scratch, so that stage-level recompute never fires, and it would also conflict with a consumer that
has already read the output incrementally.

This is the first PR in a stack that incrementally introduces pipelined shuffle dependencies and
concurrent stage scheduling.

Why are the changes needed?

Today a multi-stage job runs one stage at a time: each shuffle is fully materialized before the next
stage starts. Some workloads (initially Structured Streaming real-time mode) need the stages of a
single job to run concurrently, connected by a shuffle whose consumer reads the producer's output as
it is produced. Expressing that requires a first-class, distinguishable dependency kind that the
scheduler and shuffle layers can key off of. This PR introduces that primitive as the foundation the
subsequent scheduler and shuffle-layer changes build on, split out on its own so it can be reviewed
in isolation without any behavioral risk.

Does this PR introduce any user-facing change?

No. The new class is a @DeveloperApi addition that changes no existing behavior; nothing constructs
a PipelinedShuffleDependency yet, and it is handled identically to an ordinary ShuffleDependency
everywhere.

How was this patch tested?

New unit tests in ShuffleDependencySuite (all passing) verify that a PipelinedShuffleDependency:

  • is-a ShuffleDependency (so existing ShuffleDependency matches continue to apply) and
    preserves its fields (partitioner, key/value class names, rdd);
  • registers its shuffle through the normal path (a shuffleHandle is produced) and a second instance
    receives a distinct shuffleId;
  • keeps checksumMismatchFullRetryEnabled and checksumMismatchQueryLevelRollbackEnabled at false;
  • forwards non-default constructor arguments (aggregator, mapSideCombine) to ShuffleDependency
    correctly (positional-forwarding guard);

and that an ordinary ShuffleDependency is not a PipelinedShuffleDependency (negative case).

build/sbt 'core/testOnly org.apache.spark.shuffle.ShuffleDependencySuite'
...
Tests: succeeded 6, failed 0

Was this patch authored or co-authored using generative AI tooling?

co-authored: Claude Code (Opus 4.8)

Introduce PipelinedShuffleDependency, a ShuffleDependency subtype that declares
a shuffle's output can be read incrementally -- a consumer stage may begin
reading while the producer is still running. This is the first-class marker a
later DAGScheduler change will use to co-schedule the producer and consumer (a
'pipelined group') and to select an incremental shuffle implementation.

This PR adds only the type. On its own it behaves exactly like its parent
ShuffleDependency -- code that matches ShuffleDependency continues to treat it
as an ordinary materialized shuffle -- so it changes no existing behavior. The
concurrent-scheduling and incremental-shuffle behavior are added separately.

The parent's checksumMismatchFullRetryEnabled / checksumMismatchQueryLevelRoll-
backEnabled params are intentionally not exposed, so they stay false: their
stage-level recompute-and-rerun is moot for a pipelined group (any failure
aborts the whole group and the caller reruns from scratch), and would also
conflict with a consumer that has already read the output incrementally.

Tests: PipelinedShuffleDependency is-a ShuffleDependency, is distinguishable
from an ordinary dependency, registers its shuffle with a distinct shuffleId,
keeps the checksum retry/rollback flags false, and forwards non-default
constructor args (aggregator, mapSideCombine) to the parent correctly.

Co-authored-by: Isaac
@jerrypeng jerrypeng changed the title [WIP][SPARK-58078][CORE] Add PipelinedShuffleDependency [SPARK-58078][CORE] Add PipelinedShuffleDependency Jul 10, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant